package com.duck.inspection.config.websockt;

import com.coze.openapi.client.chat.CreateChatReq;
import com.coze.openapi.client.chat.model.ChatEvent;
import com.coze.openapi.client.chat.model.ChatEventType;
import com.coze.openapi.client.connversations.message.model.Message;
import com.coze.openapi.client.connversations.message.model.MessageObjectString;
import com.duck.inspection.config.CozeChatApiClient;
import com.duck.inspection.domain.dto.ChatStreamRequest;
import com.duck.inspection.utils.QianWenChat;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Component
@Slf4j
public class ChatWebSocketHandler extends TextWebSocketHandler {

    private final CozeChatApiClient chatApiClient;
    private final ObjectMapper objectMapper = new ObjectMapper();

    public ChatWebSocketHandler(CozeChatApiClient chatApiClient) {
        this.chatApiClient = chatApiClient;
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            ChatStreamRequest req = objectMapper.readValue(
                    message.getPayload(), ChatStreamRequest.class);

            // 1. 获取或创建 conversationId
            String convId = (String) session.getAttributes().get("conversationId");
            if (convId == null) {
                convId = chatApiClient.createConversation();
                session.getAttributes().put("conversationId", convId);
            }

            // 2. 构造文本+图像消息
            List<MessageObjectString> objs = new ArrayList<>();
            if (req.getUserMessage() != null && !req.getUserMessage().isBlank()) {
                objs.add(MessageObjectString.buildText(req.getUserMessage()));
            }
            if (req.getImageId() != null && !req.getImageId().isBlank()) {
                objs.add(MessageObjectString.buildImageByID(req.getImageId()));
            }
            if (objs.isEmpty()) {
                session.sendMessage(new TextMessage("请发送文本或图片"));
                return;
            }

            CreateChatReq chatReq = CreateChatReq.builder()
                    .botID(req.getBotID())
                    .userID(req.getUserID())
                    .conversationID(convId)
                    .messages(Collections.singletonList(
                            Message.buildUserQuestionObjects(objs)))
                    .build();


            StringBuilder fullAnswer = new StringBuilder();
            Flowable<ChatEvent> flow = chatApiClient.streamChat(chatReq);

            flow
                    .subscribeOn(Schedulers.io())
                    .filter(e -> e.getMessage() != null &&
                            (e.getEvent() == ChatEventType.CONVERSATION_MESSAGE_DELTA ||
                                    e.getEvent() == ChatEventType.CONVERSATION_CHAT_COMPLETED))
                    .map(e -> e.getMessage().getContent())
                    .subscribe(
                            fragment -> {
                                if (session.isOpen()) {
                                    try {
                                        session.sendMessage(new TextMessage(fragment));
                                    } catch (IOException ignore) { /* 可记录 */ }
                                }
                                fullAnswer.append(fragment);
                            },
                            err -> {
                                log.error("流式对话异常", err);
                                try {
                                    session.sendMessage(
                                            new TextMessage("错误: " + err.getMessage()));
                                } catch (IOException ignore) {}
                            },
                            () -> {
                                // onComplete: 生成后续建议
                                try {
                                    String sysPrompt = """
                            你是一个智能提示词生成器，当用户提问后，请严格按照要求生成：
                            1. 生成3条用户可能会再次发送的相关问题（不要反问）
                            2. 输出纯JSON数组格式，不要包含任何Markdown代码块或额外文本
                            3. 每个问题不超过15个字

                            示例格式：
                            ["问题1", "问题2", "问题3"]

                            当前用户问题：
                            """;
                                    String followUps = QianWenChat.callWithMessage(
                                            sysPrompt + req.getUserMessage()
                                    );
                                    String clean = followUps
                                            .replaceAll("```(json)?", "")
                                            .replaceAll("[\\n\\t]", "")
                                            .trim();
                                    List<String> suggestions = objectMapper.readValue(
                                            clean, new TypeReference<List<String>>() {}
                                    );
                                    String payload = objectMapper.writeValueAsString(
                                            Collections.singletonMap("suggestions", suggestions)
                                    );
                                    session.sendMessage(new TextMessage(payload));
                                } catch (Exception ex) {
                                    log.error("生成提示词失败", ex);
                                }
                            }
                    );

        } catch (Exception ex) {
            log.error("处理客户端请求失败", ex);
            try {
                session.sendMessage(
                        new TextMessage("请求处理失败: " + ex.getMessage()));
            } catch (IOException __) {
                log.error("发送失败消息异常", __);
            }
        }
    }
}
